package org.finesys.chat.core.base.client;

import com.volcengine.ark.runtime.model.completion.chat.*;
import com.volcengine.ark.runtime.model.completion.chat.ChatCompletionRequest;
import com.volcengine.ark.runtime.model.embeddings.EmbeddingRequest;
import com.volcengine.ark.runtime.model.embeddings.EmbeddingResult;
import com.volcengine.ark.runtime.model.images.generation.GenerateImagesRequest;
import com.volcengine.ark.runtime.model.images.generation.ImagesResponse;
import com.volcengine.ark.runtime.service.ArkService;
import io.reactivex.Flowable;
import org.finesys.ai.core.chat.*;
import org.finesys.ai.core.chat.ChatCompletionChoice;
import org.finesys.ai.core.completion.CompletionResponse;
import org.finesys.ai.core.embedding.Embedding;
import org.finesys.ai.core.embedding.EmbeddingResponse;
import org.finesys.ai.core.image.GenerateImagesResponse;
import org.finesys.ai.core.image.ImageData;
import org.finesys.ai.core.shared.CompletionTokensDetails;
import org.finesys.ai.core.shared.PromptTokensDetails;
import org.finesys.ai.core.shared.Usage;
import org.finesys.chat.core.base.chat.ChatRequest;
import org.finesys.chat.core.base.factory.ChatModel;
import org.finesys.chat.core.base.mapper.DoubaoConverter;
import org.springframework.util.ObjectUtils;
import org.springframework.util.StringUtils;
import reactor.core.publisher.Flux;

import java.net.URI;
import java.time.Duration;
import java.util.ArrayList;
import java.util.List;

public class DoubaoChatClient extends ChatClientAbstract {


    private ArkService arkService;

    /**
     * DoubaoChatClient 类的构造函数
     *
     * @param chatModel 聊天模型，用于初始化聊天客户端
     */
    public DoubaoChatClient(ChatModel chatModel) {
        super(chatModel);
        arkService = build(chatModel);
    }

    /**
     * 根据ChatModel构建ArkService实例
     *
     * @param chatModel 包含构建ArkService所需的配置信息的模型
     * @return 构建的ArkService实例
     * @throws IllegalArgumentException 如果apiKey为空则抛出此异常
     */
    private ArkService build(ChatModel chatModel) {
        ArkService.Builder builder = ArkService.builder();
        if (StringUtils.hasText(chatModel.getBaseUrl())) {
            builder.baseUrl(chatModel.getBaseUrl());
        }
        if (!StringUtils.hasText(chatModel.getApiKey())) {
            throw new IllegalArgumentException("apiKey is empty");
        }
        if (chatModel.getProxy() != null) {
            builder.proxy(chatModel.getProxy());
        }
        if (chatModel.getCallTimeout() > 0) {
            builder.callTimeout(Duration.ofSeconds(chatModel.getCallTimeout()));
        }
        return builder.build();
    }


    @Override
    public void updateKey(ChatModel chatModel) {
        super.updateKey(chatModel);
        if (arkService == null) {
            throw new RuntimeException("arkService is null,can not updateKey.");
        }
        //关闭旧的客户端连接
        arkService.shutdownExecutor();
        //重新构建新的DeepSeekClient实例。
        arkService = build(chatModel);
    }

    @Override
    public void register(ChatModel chatModel) {
        super.register(chatModel);
        if (arkService != null) {
            throw new RuntimeException("arkService is already initialized. Please use updateKey method to update the key.");
        }
        //重新构建新的DeepSeekClient实例。
        arkService = build(chatModel);
    }

    @Override
    public void unregister() {
        super.unregister();
        if (arkService == null) {
            throw new RuntimeException("arkService is null,can not unregister.");
        }
        //关闭旧的客户端连接
        arkService.shutdownExecutor();
    }

    /**
     * 根据ChatRequest构建ChatCompletionRequest
     *
     * @param chatRequest ChatRequest对象
     * @return 构建后的ChatCompletionRequest对象
     */
    private ChatCompletionRequest buildDoubaoChatCompletionRequest(ChatRequest chatRequest) {
        ChatCompletionRequest.Builder completionRequestBuilder = ChatCompletionRequest.builder()
                .model(chatRequest.getModel())
                .maxTokens(chatRequest.getMaxTokens())
                .temperature(chatRequest.getTemperature())
                .topP(chatRequest.getTopP())
                .n(chatRequest.getN())
                .stream(chatRequest.getStream())
                .stop(stopWords(chatRequest.getStopWords()))
                .presencePenalty(chatRequest.getPresencePenalty())
                .frequencyPenalty(chatRequest.getFrequencyPenalty())
                .logitBias(chatRequest.getLogitBias())
                .user(chatRequest.getUser());
        if (chatRequest.getMessages() != null) {
            completionRequestBuilder.messages(buildMessages(chatRequest.getMessages()));
        }

        return completionRequestBuilder.build();
    }

    private List<ChatMessage> buildMessages(List<org.finesys.chat.core.base.chat.ChatMessage> chatMessages) {
        List<ChatMessage> chatMessagesList = new ArrayList<>();
        chatMessages.forEach(chatMessage -> {
            ChatMessage.Builder chatMessageBuilder = ChatMessage.builder();
            if (StringUtils.hasText(chatMessage.getRole().name())) {
                chatMessageBuilder.role(ChatMessageRole.valueOf(chatMessage.getRole().name()));
            }
            if (!ObjectUtils.isEmpty(chatMessage.getContent())) {
                chatMessageBuilder.content(String.valueOf(chatMessage.getContent()));
            }
            if (StringUtils.hasText(chatMessage.getReasoningContent())) {
                chatMessageBuilder.reasoningContent(chatMessage.getReasoningContent());
            }

            if (!ObjectUtils.isEmpty(chatMessage.getToolCalls())) {
                chatMessageBuilder.toolCalls(DoubaoConverter.INSTANCE.convertList(chatMessage.getToolCalls()));
            }

            if (!ObjectUtils.isEmpty(chatMessage.getFunctionCall())) {
                chatMessageBuilder.functionCall(DoubaoConverter.INSTANCE.convert(chatMessage.getFunctionCall()));
            }
            chatMessagesList.add(chatMessageBuilder.build());
        });
        return chatMessagesList;
    }


    private ChatCompletionResponse buildChatCompletionResponse(ChatCompletionResult chatCompletionResult) {
        ChatCompletionResponse.Builder chatCompletionResponseBuilder = ChatCompletionResponse.builder();
        chatCompletionResponseBuilder.id(chatCompletionResult.getId());
        chatCompletionResponseBuilder.created(Long.valueOf(chatCompletionResult.getCreated()).intValue());
        chatCompletionResponseBuilder.model(chatCompletionResult.getModel());
        chatCompletionResponseBuilder.choices(buildChatCompletionChoices(chatCompletionResult.getChoices()));
        chatCompletionResponseBuilder.usage(buildUsage(chatCompletionResult.getUsage()));
        return chatCompletionResponseBuilder.build();
    }

    private List<ChatCompletionChoice> buildChatCompletionChoices(List<com.volcengine.ark.runtime.model.completion.chat.ChatCompletionChoice> chatCompletionChoices) {
        List<ChatCompletionChoice> chatCompletionChoices1 = new ArrayList<>();
        chatCompletionChoices.forEach(chatCompletionChoice -> {
            ChatCompletionChoice.Builder chatCompletionChoiceBuilder = ChatCompletionChoice.builder();
            chatCompletionChoiceBuilder.finishReason(chatCompletionChoice.getFinishReason());
            chatCompletionChoiceBuilder.index(chatCompletionChoice.getIndex());
            chatCompletionChoiceBuilder.delta(buildDelta(chatCompletionChoice.getMessage()));
            chatCompletionChoices1.add(chatCompletionChoiceBuilder.build());
        });

        return chatCompletionChoices1;
    }

    private Delta buildDelta(ChatMessage chatMessage) {
        Delta.Builder deltaBuilder = Delta.builder();
        deltaBuilder.content(chatMessage.getReasoningContent());
        if (StringUtils.hasText(chatMessage.getRole().name())) {
            deltaBuilder.role(Role.valueOf(chatMessage.getRole().name()));
        }
        if (chatMessage.getToolCalls() != null) {
            deltaBuilder.toolCalls(buildToolCalls(chatMessage.getToolCalls()));
        }
        return deltaBuilder.build();
    }

    private List<ToolCall> buildToolCalls(List<ChatToolCall> toolCalls) {
        List<ToolCall> toolCalls1 = new ArrayList<>();
        toolCalls.forEach(toolCall -> {
            ToolCall.Builder toolCallBuilder = ToolCall.builder();
            toolCallBuilder.id(toolCall.getId());
            toolCallBuilder.index(toolCall.getIndex());
            if (toolCall.getFunction() != null) {
                toolCallBuilder.function(buildFunction(toolCall.getFunction()));
            }
            if (StringUtils.hasText(toolCall.getType())) {
                toolCallBuilder.type(ToolType.valueOf(toolCall.getType()));
            }
            toolCalls1.add(toolCallBuilder.build());
        });

        return toolCalls1;
    }

    private FunctionCall buildFunction(ChatFunctionCall function) {
        FunctionCall.Builder functionCallBuilder = FunctionCall.builder();
        functionCallBuilder.name(function.getName());
        functionCallBuilder.arguments(function.getArguments());
        return functionCallBuilder.build();
    }


    private ChatCompletionResponse buildChatCompletionResponse(ChatCompletionChunk chatCompletionChunk) {
        ChatCompletionResponse.Builder chatCompletionResponseBuilder = ChatCompletionResponse.builder();
        chatCompletionResponseBuilder.id(chatCompletionChunk.getId());
        chatCompletionResponseBuilder.model(chatCompletionChunk.getModel());
        chatCompletionResponseBuilder.created(Long.valueOf(chatCompletionChunk.getCreated()).intValue());
        chatCompletionResponseBuilder.choices(buildChatCompletionChoices(chatCompletionChunk.getChoices()));
        chatCompletionResponseBuilder.usage(buildUsage(chatCompletionChunk.getUsage()));
        chatCompletionResponseBuilder.systemFingerprint(chatCompletionChunk.getObject());
        chatCompletionResponseBuilder.serviceTier(chatCompletionChunk.getServiceTier());
        return chatCompletionResponseBuilder.build();
    }

    private Usage buildUsage(com.volcengine.ark.runtime.model.Usage usage) {
        Usage.Builder usageBuilder = Usage.builder();
        usageBuilder.totalTokens(Long.valueOf(usage.getTotalTokens()).intValue());
        usageBuilder.completionTokens(Long.valueOf(usage.getCompletionTokens()).intValue());
        usageBuilder.promptTokensDetails(buildPromptTokensDetails(usage.getPromptTokensDetails()));
        usageBuilder.completionTokensDetails(buildCompletionTokensDetails(usage.getCompletionTokensDetails()));
        return usageBuilder.build();
    }

    private CompletionTokensDetails buildCompletionTokensDetails(com.volcengine.ark.runtime.model.CompletionTokensDetails completionTokensDetails) {
        CompletionTokensDetails.Builder completionTokensDetailsBuilder = CompletionTokensDetails.builder();
        completionTokensDetailsBuilder.reasoningTokens(completionTokensDetails.getReasoningTokens());
        return completionTokensDetailsBuilder.build();
    }

    private PromptTokensDetails buildPromptTokensDetails(com.volcengine.ark.runtime.model.PromptTokensDetails promptTokensDetails) {
        PromptTokensDetails.Builder promptTokensDetailsBuilder = PromptTokensDetails.builder();
        promptTokensDetailsBuilder.cachedTokens(promptTokensDetails.getCachedTokens());
        return promptTokensDetailsBuilder.build();
    }


    @Override
    public Flux<CompletionResponse> completion(ChatRequest chatRequest) {
        return super.completion(chatRequest);
    }

    @Override
    public Flux<ChatCompletionResponse> chatCompletion(ChatRequest chatRequest) {
        ChatCompletionRequest chatCompletionRequest = buildDoubaoChatCompletionRequest(chatRequest);
        //异步响应
        if (chatRequest.getStream()) {
            return Flux.create(sink -> {
                Flowable<ChatCompletionChunk> chatCompletionResult = arkService.streamChatCompletion(chatCompletionRequest);
                chatCompletionResult.subscribe(
                        chunk -> {
                            if (!sink.isCancelled()) {
                                ChatCompletionResponse chatCompletionResponse = buildChatCompletionResponse(chunk);
                                sink.next(chatCompletionResponse);
                            }
                        },
                        error -> sink.error(error),
                        () -> {
                            if (!sink.isCancelled()) {
                                sink.complete();
                            }
                        });

            });
        }
        ChatCompletionResult chatCompletionResult = arkService.createChatCompletion(chatCompletionRequest);
        ChatCompletionResponse chatCompletionResponse = buildChatCompletionResponse(chatCompletionResult);
        return Flux.just(chatCompletionResponse);
    }

    @Override
    public Flux<GenerateImagesResponse> imageGeneration(ChatRequest chatRequest) {
        ImagesResponse imagesResponse = arkService.generateImages(buildDoubaoGenerateImagesRequest(chatRequest));
        GenerateImagesResponse generateImagesResponse = buildGenerateImagesResponse(imagesResponse);
        return Flux.just(generateImagesResponse);
    }

    private GenerateImagesResponse buildGenerateImagesResponse(ImagesResponse imagesResponse) {
        GenerateImagesResponse.Builder generateImagesResponseBuilder = GenerateImagesResponse.builder();
        generateImagesResponseBuilder.data(buildImageData(imagesResponse.getData()));
        return generateImagesResponseBuilder.build();
    }

    private List<ImageData> buildImageData(List<ImagesResponse.Image> data) {
        List<ImageData> imageData = new ArrayList<>();
        data.forEach(image -> {
            ImageData.Builder imageDataBuilder = ImageData.builder();
            imageDataBuilder.url(URI.create(image.getUrl()));
            imageDataBuilder.b64Json(image.getB64Json());
            imageData.add(imageDataBuilder.build());
        });
        return imageData;
    }

    private GenerateImagesRequest buildDoubaoGenerateImagesRequest(ChatRequest chatRequest) {
        GenerateImagesRequest.Builder generateImagesRequestBuilder = GenerateImagesRequest.builder();
        generateImagesRequestBuilder.model(chatRequest.getModel());
        generateImagesRequestBuilder.prompt(chatRequest.getUserMessage());
        generateImagesRequestBuilder.size(chatRequest.getImageSize());
        generateImagesRequestBuilder.responseFormat(chatRequest.getImageResponseFormat());
        generateImagesRequestBuilder.guidanceScale(chatRequest.getGuidanceScale());
        generateImagesRequestBuilder.seed(chatRequest.getSeed());
        generateImagesRequestBuilder.watermark(chatRequest.getWatermark());
        return generateImagesRequestBuilder.build();
    }

    @Override
    public Flux<EmbeddingResponse> embedding(ChatRequest chatRequest) {
        EmbeddingResult embeddingResult = arkService.createEmbeddings(buildDoubaoEmbeddingRequest(chatRequest));
        EmbeddingResponse embeddingResponse = buildEmbeddingResponse(embeddingResult);
        return Flux.just(embeddingResponse);
    }

    private EmbeddingResponse buildEmbeddingResponse(EmbeddingResult embeddingResult) {
        EmbeddingResponse.Builder embeddingResponseBuilder = EmbeddingResponse.builder();
        embeddingResponseBuilder.model(embeddingResult.getModel());
        embeddingResponseBuilder.data(buildEmbeddingData(embeddingResult.getData()));
        embeddingResponseBuilder.usage(buildUsage(embeddingResult.getUsage()));
        return embeddingResponseBuilder.build();
    }

    private List<Embedding> buildEmbeddingData(List<com.volcengine.ark.runtime.model.embeddings.Embedding> data) {
        List<Embedding> embeddingList = new ArrayList<>(data.size());
        data.forEach(embedding -> {
            Embedding.Builder embeddingBuilder = Embedding.builder();
            embeddingBuilder.index(embedding.getIndex());
            embeddingBuilder.embedding(embedding.getEmbedding());
            embeddingList.add(embeddingBuilder.build());
        });
        return embeddingList;
    }

    private EmbeddingRequest buildDoubaoEmbeddingRequest(ChatRequest chatRequest) {
        EmbeddingRequest.Builder embeddingRequestBuilder = EmbeddingRequest.builder();
        embeddingRequestBuilder.model(chatRequest.getModel());
        embeddingRequestBuilder.input(chatRequest.getInput());
        embeddingRequestBuilder.user(chatRequest.getUser());
        return embeddingRequestBuilder.build();
    }
}
